#Copyright Okhtay Shoghi 2008

require 'uranium/hash'
require 'timeout'
require 'net/http'
require 'json/pure'
require 'test/unit'
require 'uranium/webutils'
include Test::Unit::Assertions

class Float
	alias_method :old_round, :round
	
	def round(a)
		a = (10**a)
		return (self*a).to_f.old_round.to_f/a
	end
	
	def == x
		return false if x.class != Float 
		return (self - x).abs < 0.0000000000001
	end
end

module Test
	module Unit
		class TestCase
			alias_method :old_run, :run
			@@run_once = Hash.new
					
			def run_once
				return if (@@run_once.has_key?(self.class))
					
				puts "\n" + "*"*50 + "\nRun Once:\n"
				yield
				puts "\nDone\n" + "*"*50 + "\n"
				@@run_once[self.class] = true
			end

			def run(results,&block)
				Log.test(name)
				old_run(results,&block)
			end
		end
	end
end

def load_test(args={})
	args.defaults 'thread_count'	=> 10,
				  'duration'		=> 60,
				  'iteration_count'	=> nil,
				  'fork_count'		=> 1
	
	WebUtils::reset_request_counts
	mutex = Mutex.new
	threads = []
	txn_times = []
	
	if (args['iteration_count'] == nil)
		1.upto(args['thread_count']) do |i|
			threads.push(Thread.new do				
				while(1) do
					start = Time.new
					yield
					mutex.synchronize do
						txn_times.push(Time.new-start)
					end
				end
			end)
		end
		sleep args['duration']
	else
		puts "running #{args['iteration_count']} iterations"
		threads.push(Thread.new do
			1.upto(args['iteration_count']) do |i|
				yield
			end
		end)
		threads[0].join
	end
	
	threads.each do |thread|
		thread.kill
	end
	
	require 'pp'
	c = WebUtils::get_request_counts
	c['tps'] = (txn_times.length/txn_times.sum)*args['thread_count']	#total number of txns completed by all threads / all 
	c['avg_trip_time'] = txn_times.avg
	c['calls'].each do |func,times|
		t = times.sort
		c['calls'][func] = {
			'min'			=> t[0], 
			'max'			=> t[t.length-1], 
			'tps'			=> ((t.length.to_f/t.sum)*args['thread_count'].to_f),
			'avg'			=> t.avg,
			'count'			=> t.length}
	end
	
	#pp c
	#puts "TPS: #{(c['total']['rcvd'].to_f/args['duration'].to_f).round(1)} over #{args['duration']}s"
	
	return c
end
		

def file_get_contents(path)
	File.read(path)
end

def file_put_contents(path, data)
	File.open(path, 'w') do |file| 
		file.write data
	end
end

def assert_equal_or_diff(a,b,msg="")
	require 'pp'
	if (a != b)
		pp a.diff(b) if (a.class == Hash || a.class == Array)
		
		assert_equal(a,b, msg)
	else
		#still do this to increment the assertion count
		assert_equal(a,b)
	end
end

#A twist on the assert equal function to allow tests to pass as long as the got is a superset of the expected
def assert_hash(expected, got, opts={})
	msg = opts.class == String ? opts : opts['message']
	
	if (expected.class == Hash)
		expected = [expected]
		got = [got]
	end
	
	#		if (opts.class == Hash && opts['allow_superset'])
		0.upto(expected.length-1) do |i|
			keys = expected[i].keys
			compare_to = {}
			keys.each do |key|
				compare_to[key] = got[i][key]
			end
			
			diff_strs = []
			(got[i].keys - keys).each do |k|
				diff_strs.push("key #{k}: #{got[i][k].inspect}")
			end
			

			Log.debug "\n\n*** Response contained extra keys: \n#{diff_strs.join("\n")}\n\n"
			assert_equal(expected[i], compare_to, diff(expected[i], compare_to))
		end
	#		else
	#		assert_equal(expected, got, msg)
	#end
end

def diff(a,b)
	if (a.class != b.class)
		Log.debug "Cannot diff #{a.class} to #{b.class}"
		return
	end
	
	if (a.class == Array)
		if (a.length != b.length)
			Log.debug "Arrays have different lengths"
			return
		end
	
		diffs = []
		0.upto(a.length-1) do |i|
			diff = hash_diff(a[i],b[i])
			if (diff.length > 0)
				diffs.push("index #{i}:------------------------------------------------------")
				diffs += diff
			end
		end
		return diffs.join("\n")
	elsif (a.class == Hash)
		return hash_diff(a,b).join("\n")
	end
end

def hash_diff(hash1, hash2)
	diffs = []
	(hash1.keys + hash2.keys).uniq.each do |key|
		if (hash1[key] != hash2[key])
			diffs.push("#{key}: #{hash1[key]} vs #{hash2[key]}")
		end
	end
	return diffs
end

def assert_in_list(item, list)
	full_message = "#{item.inspect} not in list #{list.inspect}"
	
	assert_block(full_message) {
		list.index(item) 
	}
end

def assert_false(assertion, message=nil)
	assert(!assertion, message)
end

def spawn_process(cmd,sock=nil)
	if is_windows?
		require 'win32ole'
		
		sh = WIN32OLE.new('Wscript.Shell')
		sh.run(cmd)
	else
		pid = Process.fork{
			puts "Starting: #{cmd}"
			
			sock.close if sock!=nil
			exec(cmd)
		}
		Process.detach(pid)
	end
end

def find_first_open_port
	0.upto(10) do |i|
		puts "trying #{i}"
		return (4444+i) if is_port_open?(4444+i)
	end
end

def is_port_open?(port,host='0.0.0.0')
	require 'socket'
	require 'timeout'

	begin
		Timeout::timeout(1) do
			begin
				s = TCPServer.open('0.0.0.0', port)
				s.close
				puts "open"
				return false
			rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH
				return true
			end
		end
	rescue Timeout::Error
	end

	puts "timeout"
	return true
end

def kill(pid)
	if is_windows?
		`taskkill pid`
	else
		`kill -9 #{pid}`
	end
end

def get_browser_versions
	plat			= RUBY_PLATFORM.match(/darwin|win/).to_a[0]
	versions	= Hash.new
	
	case plat
		when 'win32','win64','mingw32','mingw64'
			begin
			require 'win32/registry'
			
			if (reg =	Win32::Registry::HKEY_LOCAL_MACHINE.open('Software\\Mozilla\\Mozilla Firefox\\'))
				versions['Firefox'] = reg.read('CurrentVersion')[1].match(/(.*?) /).to_a[1]
			end

			versions['IE'] = Win32::Registry::HKEY_LOCAL_MACHINE.open('Software\\Microsoft\\Internet Explorer\\').read('Version')[1]
			rescue Exception => e
				puts e.inspect
			end
		when 'darwin'
			paths = {
				'Firefox' => "/Applications/Firefox.app",
				'Safari'	=> "/Applications/Safari.app",
				'Opera'		=> "/Applications/Opera.app",
			}
			
			paths.each do |browser, path|
				if File.exist? path
					versions[browser] = File.read("#{paths[browser]}/Contents/info.plist").match(/ShortVersion.*?<string>(.*?)<\/string>/m).to_a[1]
				end
			end
	end

	return versions
end

if (!defined?(SeleniumCommandError))
	class SeleniumCommandError < RuntimeError 
	end
end

def wait_for_success(timeout=15, retry_wait=0.2)
	e = nil
	start = Time.new.to_i
 
	begin
		Timeout::timeout(timeout) do
			while(true)
				begin
					yield
				rescue Test::Unit::AssertionFailedError => e
				rescue SeleniumCommandError => e
					#rescue until either timeout or success
				else
					return
				end
	 
				if (Time.new.to_i - start >= timeout)
						raise(e)
				end
				sleep retry_wait #sleep between retries
			end
		end
	rescue Timeout::Error => et
		raise (e ? e : et) #if we timeout show original error not the timeout 
	end
end

def is_windows?
	return RUBY_PLATFORM.match(/win32|win64|mingw32|mingw64/)
end

def fromJSON(str)
	return JSON.parse(str)
end

def toJSON(obj)
	json = obj.inspect.gsub(/(^|([^\\]"|[\]]|=>\d+), ):[a-z]\w+\=\>/i) { |m| m.sub(/:(\w+)/) { |m| "\"#{m[1..-1]}\"" } }
	json.gsub(/(?:=>)(?=("{1}|[\[](?=("{1}|[\[]|\d+, \"))|\d+(, \"|\})))/, ": ")
end
